import torch
import torch.nn as nn
import torch.nn.functional as F
from collections import OrderedDict


class _DenseLayer(nn.Sequential):
    def __init__(self, num_input_features, growth_rate, bn_size, drop_rate):
        super(_DenseLayer, self).__init__()

        self.add_module('norm.1', nn.BatchNorm2d(num_input_features))
        self.add_module('relu.1', nn.ReLU(inplace=True))
        self.add_module('conv.1', nn.Conv2d(num_input_features, bn_size*growth_rate,
                                            kernel_size=1, stride=1, bias=False))

        self.add_module('norm.2', nn.BatchNorm2d(bn_size*growth_rate))
        self.add_module('relu.2', nn.ReLU(inplace=True))
        self.add_module('conv.2', nn.Conv2d(bn_size*growth_rate, growth_rate,
                                            kernel_size=3, stride=1, padding=1, bias=False))

        self.dropout = nn.Dropout(drop_rate)

    def forward(self, x):
        new_features = super(_DenseLayer, self).forward(x)
        new_features = self.dropout(new_features)
        return torch.cat([x, new_features], 1)


class _DenseBlock(nn.Sequential):
    def __init__(self, num_layers, num_input_features, bn_size, growth_rate, drop_rate):
        super(_DenseBlock, self).__init__()
        for i in range(num_layers):
            layer = _DenseLayer(
                num_input_features + i*growth_rate,
                growth_rate, bn_size, drop_rate
            )
            self.add_module('denselayer%d' % (i + 1), layer)


class _Transition(nn.Sequential):
    def __init__(self, num_input_features, num_output_features):
        super(_Transition, self).__init__()
        self.add_module('norm', nn.BatchNorm2d(num_input_features))
        self.add_module('relu', nn.ReLU(inplace=True))
        self.add_module('conv', nn.Conv2d(num_input_features, num_output_features,
                                          kernel_size=1, stride=1, bias=False))
        self.add_module('pool', nn.AvgPool2d(kernel_size=2, stride=2))


class DenseNet(nn.Module):
    """Densenet-BC model class, based on
    `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`

    Arguments:
        growth_rate (int) - how many filters to add each layer (`k` in paper).
        block_config (list of 4 ints) - how many layers in each pooling block.
        num_init_features (int) - the number of filters to learn
            in the first convolution layer.
        bn_size (int) - multiplicative factor for number of bottleneck layers
            (i.e. bn_size * k features in the bottleneck layer).
        drop_rate (float) - dropout rate after each dense layer.
        final_drop_rate (float) - dropout rate before final fc layer.
        num_classes (int) - number of classification classes.
    """
    def __init__(self, growth_rate=12, block_config=(8, 12, 10),
                 num_init_features=48, bn_size=4, drop_rate=0.25,
                 final_drop_rate=0.25, num_classes=200):

        super(DenseNet, self).__init__()

        # first convolution
        self.features = nn.Sequential(OrderedDict([
            ('conv0', nn.Conv2d(3, num_init_features,
                                kernel_size=3, stride=1, padding=1, bias=False)),
            ('norm0', nn.BatchNorm2d(num_init_features)),
            ('relu0', nn.ReLU(inplace=True)),
            ('pool0', nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
        ]))

        # each denseblock
        num_features = num_init_features
        for i, num_layers in enumerate(block_config):
            block = _DenseBlock(
                num_layers=num_layers, num_input_features=num_features,
                bn_size=bn_size, growth_rate=growth_rate, drop_rate=drop_rate
            )
            self.features.add_module('denseblock%d' % (i + 1), block)
            num_features = num_features + num_layers*growth_rate
            if i != len(block_config) - 1:
                trans = _Transition(
                    num_input_features=num_features,
                    num_output_features=num_features // 2
                )
                self.features.add_module('transition%d' % (i + 1), trans)
                num_features = num_features // 2

        # final batch norm
        self.features.add_module('norm4', nn.BatchNorm2d(num_features))

        # linear layer
        self.relu = nn.ReLU(inplace=True)
        self.avg_pool = nn.AvgPool2d(kernel_size=7)
        self.dropout = nn.Dropout(p=final_drop_rate)
        self.classifier = nn.Linear(num_features, num_classes)

    def forward(self, x):
        features = self.features(x)
        out = self.relu(features)
        out = self.avg_pool(out).view(features.size(0), -1)
        out = self.dropout(out)
        logits = self.classifier(out)
        return logits
